from manim import *
# manim -pql a01\SieveOfEratosthenesScene.py SieveOfEratosthenesScene
class SieveOfEratosthenesScene(Scene):
    def __init__(self, upper_limit=100, **kwargs):
        self.upper_limit = upper_limit
        super().__init__(**kwargs)

    def construct(self):
        # 创建一个1到upper_limit的数字网格
        numbers = VGroup()
        for i in range(1, self.upper_limit + 1):
            number = Text(str(i), font_size=24)
            numbers.add(number)
        
        # 计算网格的行和列数
        rows = cols = 10
        numbers.arrange_in_grid(rows=rows, cols=cols, buff=0.3)
        self.play(Write(numbers))
        self.wait(1)

        # 添加标题
        title = Text("埃拉托色尼筛法", font_size=36).to_edge(UP)
        self.play(Write(title))
        self.wait(1)

        # 记录已经筛出的倍数
        sieved_multiples = set()

        # 实现埃拉托色尼筛法
        primes = []
        for i in range(2, self.upper_limit + 1):
            if not any(i % p == 0 for p in primes):
                primes.append(i)
                # 高亮显示当前的质数
                prime_index = i - 1
                prime_number = numbers[prime_index]
                self.play(prime_number.animate.scale(1.5).set_color(YELLOW))

                # 筛选掉倍数
                multiples = []
                new_multiples = []
                for j in range(i*2, self.upper_limit + 1, i):
                    if j not in sieved_multiples:
                        multiple_index = j - 1
                        multiple_number = numbers[multiple_index]
                        multiples.append(multiple_number)
                        new_multiples.append(j)
                        sieved_multiples.add(j)
                
                if multiples:
                    # 添加说明文字
                    explanation = Text(f"{i} 是质数, 筛出 {i} 的倍数", font_size=24).next_to(numbers, DOWN)
                    self.play(Write(explanation))
                    self.wait(1)

                    # 为倍数添加外边框
                    borders = VGroup(*[SurroundingRectangle(multiple, color=RED) for multiple in multiples])
                    self.play(Write(borders), run_time=0.5)
                    self.wait(0.5)
                    # 把数字的颜色变成灰色,透明度降低

                    self.play(*[multiple.animate.set_color(GRAY).set_opacity(0.5) for multiple in multiples], run_time=0.5)
                    self.play(FadeOut(borders), run_time=0.5)
                    self.wait(0.5)
                    self.play(FadeOut(explanation))

                else:
                    # 如果所有倍数都已经筛出，给出提示
                    explanation = Text(f"{i} 是质数, 所有倍数已筛出", font_size=24).next_to(numbers, DOWN)
                    self.play(Write(explanation))
                    self.wait(1)
                    self.play(FadeOut(explanation))
                if i == 17:
                    # 高亮显示剩余所有质数
                    remaining_primes = [numbers[n-1] for n in range(i+1, self.upper_limit + 1) if n not in sieved_multiples]
                    self.play(*[prime.animate.scale(1.5).set_color(YELLOW) for prime in remaining_primes])
                    break                
                self.wait(0.5)
        
        # 最后显示所有质数
        self.play(*[number.animate.set_color(BLUE) for number in numbers if number.get_color() == YELLOW])
        self.wait(2)

# 运行命令
# manim -pql your_script.py SieveOfEratosthenesScene --upper_limit 50